package com.cc.wechat.websocket.handler;

import com.cc.wechat.websocket.message.WebSocketMessage;
import com.cc.wechat.websocket.message.WebSocketMessageType;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket处理器
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class WebSocketHandler extends TextWebSocketHandler {

    private final ObjectMapper objectMapper;
    
    // 用户会话管理，key为用户ID，value为用户的WebSocket会话
    private static final ConcurrentHashMap<Long, WebSocketSession> USER_SESSIONS = new ConcurrentHashMap<>();

    /**
     * 当WebSocket连接建立时调用
     * @param session WebSocket会话
     * @throws Exception 异常
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        Long userId = (Long) session.getAttributes().get("userId");
        if (userId != null) {
            // 存储用户会话
            USER_SESSIONS.put(userId, session);
            log.info("用户{}建立WebSocket连接", userId);
        }
    }

    /**
     * 处理WebSocket消息
     * @param session WebSocket会话
     * @param message WebSocket消息
     * @throws Exception 异常
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        Long userId = (Long) session.getAttributes().get("userId");
        
        try {
            // 解析消息
            WebSocketMessage<?> wsMessage = objectMapper.readValue(payload, WebSocketMessage.class);
            // 处理消息
            handleMessage(userId, wsMessage);
        } catch (Exception e) {
            log.error("处理WebSocket消息异常", e);
            sendErrorMessage(session, "消息处理失败");
        }
    }

    /**
     * 当WebSocket连接关闭时调用
     * @param session WebSocket会话
     * @param status 关闭状态
     * @throws Exception 异常
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        Long userId = (Long) session.getAttributes().get("userId");
        if (userId != null) {
            // 移除用户会话
            USER_SESSIONS.remove(userId);
            log.info("用户{}断开WebSocket连接", userId);
        }
    }

    /**
     * 发送消息给指定用户
     * @param userId 用户ID
     * @param message WebSocket消息
     */
    public void sendMessage(Long userId, WebSocketMessage<?> message) {
        WebSocketSession session = USER_SESSIONS.get(userId);
        if (session != null && session.isOpen()) {
            try {
                String messageText = objectMapper.writeValueAsString(message);
                session.sendMessage(new TextMessage(messageText));
            } catch (IOException e) {
                log.error("发送WebSocket消息异常", e);
            }
        }
    }

    /**
     * 发送错误消息
     * @param session WebSocket会话
     * @param errorMessage 错误消息
     * @throws IOException 异常
     */
    private void sendErrorMessage(WebSocketSession session, String errorMessage) throws IOException {
        WebSocketMessage<String> errorWsMessage = WebSocketMessage.<String>builder()
                .type(WebSocketMessageType.ERROR)
                .data(errorMessage)
                .timestamp(System.currentTimeMillis())
                .build();
        
        session.sendMessage(new TextMessage(objectMapper.writeValueAsString(errorWsMessage)));
    }

    /**
     * 处理接收到的消息
     * @param userId 用户ID
     * @param message WebSocket消息
     */
    private void handleMessage(Long userId, WebSocketMessage<?> message) {
        // 根据消息类型处理不同的消息
        switch (message.getType()) {
            case PRIVATE_MSG:
                handlePrivateMessage(message);
                break;
            case GROUP_MSG:
                handleGroupMessage(message);
                break;
            case HEARTBEAT:
                handleHeartbeat(userId);
                break;
            default:
                log.warn("未知的消息类型: {}", message.getType());
        }
    }

    /**
     * 处理私聊消息
     * @param message WebSocket消息
     */
    private void handlePrivateMessage(WebSocketMessage<?> message) {
        // 发送消息给接收者
        if (message.getReceiverId() != null) {
            sendMessage(message.getReceiverId(), message);
        }
    }

    /**
     * 处理群聊消息
     * @param message WebSocket消息
     */
    private void handleGroupMessage(WebSocketMessage<?> message) {
        // TODO: 获取群成员并发送消息
    }

    /**
     * 处理心跳消息
     * @param userId 用户ID
     */
    private void handleHeartbeat(Long userId) {
        // 更新用户在线状态
        WebSocketMessage<String> heartbeatResponse = WebSocketMessage.<String>builder()
                .type(WebSocketMessageType.HEARTBEAT)
                .data("pong")
                .timestamp(System.currentTimeMillis())
                .build();
        
        sendMessage(userId, heartbeatResponse);
    }

    /**
     * 判断用户是否在线
     * @param userId 用户ID
     * @return 是否在线
     */
    public boolean isUserOnline(Long userId) {
        WebSocketSession session = USER_SESSIONS.get(userId);
        return session != null && session.isOpen();
    }
} 